"""
@description: Utils of face detection, face recognition, and training.
@author: Zzay
@create: 2022/07/05
"""
import face_recognition
import cv2
import pickle
import os


class FaceUtils:
    """
    Utils class about facial processing.
    """
    # set detection method ("hog" or "cnn, default is "hog")
    detection_method = "hog"
    tolerance = 0.3

    def __init__(self, encoding_file_path=None):
        """
        Can be set a model file that we've trained out.
        :param encoding_file_path: Output file of training model
        """
        if encoding_file_path:
            self.load_embeddings(encoding_file_path)

    def load_embeddings(self, encoding_file_path):
        """
        Load the known embeddings.
        :param encoding_file_path: Path of the encoding file
        :return: Load the embeddings
        """
        # load the known faces and embeddings
        print("[INFO] Loading face encodings...")
        self.data = pickle.loads(open(encoding_file_path, "rb").read())

    def get_face_location(self, image):
        """
        Face detection.
        :param image: Input image
        :return: A list of face location
        """
        # convert to gray image for less cost on processing
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # detect face
        face_location_list = []
        face_locations = face_recognition.face_locations(
            img=gray,
            number_of_times_to_upsample=1,
            model=self.detection_method
        )
        # append to list
        for (top, right, bottom, left) in face_locations:
            face_location_list.append((left, top, right, bottom))
        return face_location_list

    def get_face_location_and_name(self, image):
        """
        Face recognition.
        :param image: Input image
        :return: A list of face location and the name of the person
        """
        # convert the input frame from BGR to RGB
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # compute the facial embeddings for each face
        boxes = face_recognition.face_locations(
            img=rgb,
            model=self.detection_method
        )
        encodings = face_recognition.face_encodings(rgb, boxes)

        # initialize the list of names for each face detected
        # loop over the facial embeddings
        names = []
        for encoding in encodings:
            # attempt to match faces to our known encodings
            matches = face_recognition.compare_faces(
                known_face_encodings=self.data["encodings"],
                face_encoding_to_check=encoding,
                tolerance=self.tolerance
            )
            name = "Unknown"
            # check to see if we have found a match
            if True in matches:
                # find the indices of all matched faces
                # initialize a dict to count the total number of times each face was matched
                matched_idxs = [i for (i, b) in enumerate(matches) if b]
                # loop over the matched indices and maintain a count for each recognized faces
                counts = {}
                for i in matched_idxs:
                    name = self.data["names"][i]
                    counts[name] = counts.get(name, 0) + 1
                # determine the recognized face with the largest number of votes
                name = max(counts, key=counts.get)
            # update the list of names
            names.append(name)

        face_location_list = []
        for ((top, right, bottom, left)) in boxes:
            face_location_list.append((left, top, right, bottom))
        return face_location_list, names

    def save_embeddings(self, image_paths, output_encoding_file_path):
        """
        It's used for training the model of face recognition, and save the model to the disk.
        :param image_paths: Paths of input images for training
        :param output_encoding_file_path: Output path of the encoding file
        :return: Train the model of face recognition, and save the model
        """
        # initialize the list of known encodings and known names
        known_encodings = []
        known_names = []

        # loop over the image paths
        warning = ""
        for (i, image_path) in enumerate(image_paths):
            # extract the person name from the image path
            print("[INFO] processing image {}/{}".format(i + 1, len(image_paths)))
            name = image_path.split(os.path.sep)[-2]  # person name
            # load the input image and convert it from RGB(OpenCV ordering) to RGB(dlib ordering)
            image = cv2.imread(image_path)
            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            # detect the coordinates of the bounding boxes
            boxes = face_recognition.face_locations(
                img=rgb,
                model=self.detection_method
            )
            # compute the facial embedding for the face
            encodings = face_recognition.face_encodings(rgb, boxes)
            if len(encodings) != 1:
                os.remove(image_path)
                warning += "[WARNING] Detect %d faces in %s.This file is deleted.\n" % (len(encodings), image_path)
                continue
            # loop over the encodings
            for encoding in encodings:
                # add each encoding + name to the set of known names
                known_encodings.append(encoding)
                known_names.append(name)

        # dump the facial encodings + names to disk
        print("[INFO] Serializing encodings...")
        data = {
            "encodings": known_encodings,
            "names": known_names
        }
        f = open(output_encoding_file_path, "wb")
        f.write(pickle.dumps(data))
        f.close()

        if warning:
            print(warning)
